﻿Option Strict On

Imports Alphaleonis.Win32.Vss
Imports Alphaleonis.Win32.Filesystem
Imports Sunts.HyperV.Common

Public Enum BackupEvent
    StartingBackupRoutine
    SettingBackupDevice
    DeviceOffline
    BackupExpired
    InitializingVSS
    EmptyMetadata
    StartingSnaphotSet
    SnapshotSetDone
    CopyingFile
    FileCopyError
    FileCopyComplete
    DeletingSnapshotSet
    Cancel
    Complete
    JobReschedule
    ErrorSendingMail
    Exception
End Enum

Public Enum BackupCompletionStatus
    Success
    DeviceError
    VSSError
    Timeout
    Cancelled
    CopyError
    Unknown
End Enum

Public Class BackupProgressEventArgs
    Inherits EventArgs
    Public Property FileLocation As String
    Public Property FileDestination As String
    Public Property PercentCopied As Int32
    Public Property Cancel As Boolean
    Public Property BackupEvent As BackupEvent
    Public Property VssComponentCollection As VssComponentCollection
    Public Property SnapshotVolumeCollection As SnapshotVolumeCollection
    Public Property Exception As Exception
    Public Property BackupCompletionStatus As BackupCompletionStatus
End Class

Public Class Backup
    Implements IDisposable

    Public Event BackupProgress As EventHandler(Of BackupProgressEventArgs)
    Public Property BackupProgressEventArgs As BackupProgressEventArgs
    Public Property BackupSchedule As BackupSchedule
    Public Property LogName As String
    Private Property ConnectionString As String
    Private Disposed As Boolean

    Private Function CopyProgress(ByVal totalFileSize As Long, ByVal totalBytesTransferred As Long, ByVal streamSize As Long, ByVal streamBytesTransferred As Long, ByVal dwStreamNumber As UInteger, ByVal dwCallbackReason As CopyProgressCallbackReason, ByVal UserData As Object) As CopyProgressResult
        Dim BackupProgressEventArgs As BackupProgressEventArgs = TryCast(UserData, BackupProgressEventArgs)
        Dim NewPercent As Int32 = 0

        If BackupProgressEventArgs.Cancel = True Then
            WriteToLogFile("CopyProgress: File copy has been cancelled", LogName)
            BackupProgressEventArgs = Nothing
            Return CopyProgressResult.Cancel
        End If

        NewPercent = Convert.ToInt32(totalBytesTransferred / totalFileSize * 100)

        If BackupProgressEventArgs.PercentCopied <> NewPercent Then
            Console.WriteLine(BackupProgressEventArgs.PercentCopied & " " & NewPercent)
            BackupProgressEventArgs.PercentCopied = NewPercent
            Me.BackupSchedule.BackupJobSchedule.Progress = "Copying file: " & Path.GetFileName(BackupProgressEventArgs.FileDestination) & " (" & BackupProgressEventArgs.PercentCopied & "%)"
            Me.BackupSchedule.BackupJobSchedule.Update(ConnectionString)
        End If

        BackupProgressEventArgs = Nothing
        Return CopyProgressResult.Continue

    End Function

    Public Sub New(ByVal ConnectionString As String)
        BackupProgressEventArgs = New BackupProgressEventArgs
        AddHandler BackupProgress, AddressOf ProgressManager
        Me.ConnectionString = ConnectionString
    End Sub

    Public Function ProcessBackup() As Boolean
        Dim oVSSImpl As IVssImplementation
        Dim HyperVWriterGuid As Guid = New Guid("66841cd4-6ded-4f4b-8f17-fd23f8ddc3de")
        Dim Metadata As IVssExamineWriterMetadata
        Dim VssGuid As Guid
        Dim SnapshotVolumeCollection As SnapshotVolumeCollection
        Dim VssComponentCollection As VssComponentCollection
        Dim SnapshotVolume As SnapshotVolume
        Dim BackupItem As BackupItem
        Dim BackupItemCollection As BackupItemCollection


        Try
            BackupItemCollection = New BackupItemCollection
            SnapshotVolumeCollection = New SnapshotVolumeCollection
            VssComponentCollection = New VssComponentCollection

            Using oVSS As IVssBackupComponents = VssUtils.LoadImplementation().CreateVssBackupComponents()
                Call RaiseBackupEvent(BackupEvent.InitializingVSS)

                oVSS.InitializeForBackup(Nothing)
                oVSS.SetBackupState(True, True, VssBackupType.Full, False)
                oVSS.SetContext(VssSnapshotContext.Backup)
                oVSS.EnableWriterClasses(New Guid() {HyperVWriterGuid})
                oVSS.GatherWriterMetadata()

                Metadata = oVSS.WriterMetadata.Where(Function(o) o.WriterId.Equals(HyperVWriterGuid)).FirstOrDefault

                If Metadata Is Nothing Then
                    Call RaiseBackupEvent(BackupEvent.EmptyMetadata)
                    Return False
                End If

                For Each component As IVssWMComponent In Metadata.Components
                    If component.ComponentName.ToUpper = BackupSchedule.VirtualMachine.VirtualMachineGuid.ToString.ToUpper Then
                        oVSS.AddComponent(Metadata.InstanceId, Metadata.WriterId, component.Type, component.LogicalPath, component.ComponentName)
                        VssComponentCollection.Add(component)
                    End If
                Next

                BackupProgressEventArgs.VssComponentCollection = VssComponentCollection

                If VssComponentCollection.Count > 0 Then
                    VssGuid = oVSS.StartSnapshotSet()

                    For Each component As IVssWMComponent In VssComponentCollection
                        For Each File As VssWMFileDescriptor In component.Files
                            If Not File.IsRecursive Then
                                BackupItem = New BackupItem(File.Path, File.FileSpecification)

                                If Not SnapshotVolumeCollection.ContainsVolume(BackupItem.Volume) Then
                                    SnapshotVolume = New SnapshotVolume
                                    SnapshotVolume.Volume = BackupItem.Volume
                                    SnapshotVolume.SnapshotGuid = oVSS.AddToSnapshotSet(BackupItem.Volume, Guid.Empty)
                                    SnapshotVolumeCollection.Add(SnapshotVolume)
                                    BackupItem.SnapshotGuid = SnapshotVolume.SnapshotGuid
                                Else
                                    BackupItem.SnapshotGuid = SnapshotVolumeCollection.GetSnapshotVolume(BackupItem.Volume).SnapshotGuid
                                End If
                                BackupItemCollection.Add(BackupItem)
                            End If
                        Next
                    Next

                    BackupProgressEventArgs.SnapshotVolumeCollection = SnapshotVolumeCollection

                    oVSS.PrepareForBackup()
                    Call RaiseBackupEvent(BackupEvent.StartingSnaphotSet)
                    oVSS.DoSnapshotSet()
                    Call RaiseBackupEvent(BackupEvent.SnapshotSetDone)

                    For Each SnapshotVolume In SnapshotVolumeCollection
                        SnapshotVolume.SnapshotDeviceObject = oVSS.GetSnapshotProperties(SnapshotVolume.SnapshotGuid).SnapshotDeviceObject
                    Next

                    For Each BackupItem In BackupItemCollection
                        BackupItem.SnapshotDeviceObject = SnapshotVolumeCollection.GetSnapshotVolume(BackupItem.Volume).SnapshotDeviceObject
                    Next

                    Try
                        BackupFiles(BackupItemCollection)
                    Catch ex As Exception
                        BackupProgressEventArgs.Exception = ex
                        Call RaiseBackupEvent(BackupEvent.FileCopyError)
                    End Try

                    For Each component As IVssWMComponent In VssComponentCollection
                        oVSS.SetBackupSucceeded(Metadata.InstanceId, Metadata.WriterId, component.Type, component.LogicalPath, component.ComponentName, True)
                    Next

                    oVSS.BackupComplete()
                    Call RaiseBackupEvent(BackupEvent.DeletingSnapshotSet)
                    oVSS.DeleteSnapshotSet(VssGuid, True)
                    VssGuid = Guid.Empty
                End If

            End Using

            If Me.BackupProgressEventArgs.BackupCompletionStatus = Nothing Then
                Me.BackupProgressEventArgs.BackupCompletionStatus = BackupCompletionStatus.Success
            End If
          

            Return True
        Catch ex As Exception
            BackupProgressEventArgs.Exception = ex
            Call RaiseBackupEvent(BackupEvent.Exception)
            BackupProgressEventArgs.BackupCompletionStatus = BackupCompletionStatus.VSSError
            Throw ex
        Finally
            oVSSImpl = Nothing
            HyperVWriterGuid = Nothing
            Metadata = Nothing
            VssGuid = Nothing
            SnapshotVolumeCollection = Nothing
            VssComponentCollection = Nothing
            SnapshotVolume = Nothing
            BackupItem = Nothing
            BackupItemCollection = Nothing
        End Try
    End Function

  
    Private Sub BackupFiles(ByVal BackupItemCollection As BackupItemCollection)
        Dim CopyProgressRoutine As CopyProgressRoutine

        Try
            For Each BackupItem As BackupItem In BackupItemCollection
                If BackupProgressEventArgs.Cancel = True Then
                    Exit For
                End If

                BackupProgressEventArgs = New BackupProgressEventArgs
                CopyProgressRoutine = New CopyProgressRoutine(AddressOf CopyProgress)

                BackupProgressEventArgs.FileLocation = IO.Path.Combine(BackupItem.SnapshotDeviceObject, BackupItem.PathFromVolume, BackupItem.FileName)
                BackupProgressEventArgs.FileDestination = IO.Path.Combine(BackupSchedule.BackupFolder, BackupItem.PathFromVolume, BackupItem.FileName)

                IO.Directory.CreateDirectory(IO.Path.Combine(BackupSchedule.BackupFolder, BackupItem.PathFromVolume))

                Call RaiseBackupEvent(BackupEvent.CopyingFile)
                Alphaleonis.Win32.Filesystem.File.Copy(BackupProgressEventArgs.FileLocation, BackupProgressEventArgs.FileDestination, CopyOptions.None, CopyProgressRoutine, BackupProgressEventArgs, True)
                Call RaiseBackupEvent(BackupEvent.FileCopyComplete)

            Next
        Catch ex As Exception
            Throw New Exception("Backup:BackupFiles: " & ex.Message)
        Finally
            CopyProgressRoutine = Nothing
        End Try

    End Sub

    Public Sub ProgressManager(ByVal sender As Object, ByVal BackupProgressEventArgs As BackupProgressEventArgs)
        If Me.BackupSchedule.BackupJobSchedule IsNot Nothing Then
            Select Case BackupProgressEventArgs.BackupEvent
                Case BackupEvent.SettingBackupDevice
                    WriteToLogFile("Using " & Me.BackupSchedule.BackupDevice.DevicePath & " for the backup", LogName)
                    Exit Select
                Case BackupEvent.StartingBackupRoutine
                    WriteToLogFile("Backup schedule started for machine " & Me.BackupSchedule.VirtualMachine.VirtualMachineName, LogName)
                    Exit Select
                Case BackupEvent.BackupExpired
                    Me.BackupProgressEventArgs.BackupCompletionStatus = BackupCompletionStatus.Timeout
                    WriteToLogFile("The backup for machine " & Me.BackupSchedule.VirtualMachine.VirtualMachineName & " has been cancelled as the job has timed out", LogName)
                    Exit Select
                Case BackupEvent.DeviceOffline
                    Me.BackupProgressEventArgs.BackupCompletionStatus = BackupCompletionStatus.DeviceError
                    WriteToLogFile("The backup for machine " & Me.BackupSchedule.VirtualMachine.VirtualMachineName & " has been cancelled as the backup device could not be written to", LogName)
                    Exit Select
                Case BackupEvent.InitializingVSS
                    WriteToLogFile("Initializing VSS", LogName)
                    Exit Select
                Case BackupEvent.EmptyMetadata
                    WriteToLogFile("The VSS writer metadata was empty", LogName)
                    Exit Select
                Case BackupEvent.StartingSnaphotSet
                    BackupSchedule.BackupJobSchedule.Progress = "Processing Snapshots"
                    BackupSchedule.BackupJobSchedule.Update(ConnectionString)
                    WriteToLogFile("Starting snapshot set for:", LogName)
                    For Each Component As IVssWMComponent In BackupProgressEventArgs.VssComponentCollection
                        WriteToLogFile(Component.ComponentName, LogName)
                    Next
                    WriteToLogFile("Volumes:", LogName)
                    For Each Volume As SnapshotVolume In BackupProgressEventArgs.SnapshotVolumeCollection
                        WriteToLogFile(Volume.Volume, LogName)
                    Next
                    Exit Select
                Case BackupEvent.CopyingFile
                    WriteToLogFile("Backing up file: " & BackupProgressEventArgs.FileLocation, LogName)
                    WriteToLogFile("Backing up file to: " & BackupProgressEventArgs.FileDestination, LogName)
                    Exit Select
                Case BackupEvent.FileCopyError
                    Me.BackupProgressEventArgs.BackupCompletionStatus = BackupCompletionStatus.CopyError
                    WriteToLogFile("Error copying file: " & BackupProgressEventArgs.Exception.Message, LogName)
                    Exit Select
                Case BackupEvent.FileCopyComplete
                    WriteToLogFile("Backup complete for file: " & BackupProgressEventArgs.FileLocation, LogName)
                    Exit Select
                Case BackupEvent.DeletingSnapshotSet
                    BackupSchedule.BackupJobSchedule.Progress = "Cleaning up snapshots"
                    BackupSchedule.BackupJobSchedule.Update(ConnectionString)
                    WriteToLogFile("Deleting snapshot set", LogName)
                    Exit Select
                Case BackupEvent.Complete
                    WriteToLogFile("Backup job for " & Me.BackupSchedule.VirtualMachine.VirtualMachineName & " has completed", LogName)
                    Exit Select
                Case BackupEvent.Cancel
                    WriteToLogFile("Cancelling backup job", LogName)
                    Me.BackupProgressEventArgs.BackupCompletionStatus = BackupCompletionStatus.Cancelled
                    Me.BackupSchedule.BackupJobSchedule.Cancelled = True
                    Exit Select
                Case BackupEvent.JobReschedule
                    WriteToLogFile("Add new backup job to the schedule: " & Me.BackupSchedule.BackupJobSchedule.ScheduleDateTime, LogName)
                    Exit Select
                Case BackupEvent.ErrorSendingMail
                    WriteToLogFile("Could not send email to recipients: " & BackupProgressEventArgs.Exception.Message, LogName)
                    Exit Select
                Case BackupEvent.Exception
                    WriteToLogFile("Error: " & BackupProgressEventArgs.Exception.Message, LogName)
                    Me.BackupSchedule.BackupJobSchedule.Cancelled = True
                    Exit Select

            End Select
           BackupProgressEventArgs.Cancel = Me.BackupSchedule.BackupJobSchedule.Cancelled
        End If
    End Sub

    Public Sub RaiseBackupEvent(ByVal BackupEvent As BackupEvent)
        BackupProgressEventArgs.BackupEvent = BackupEvent
        RaiseEvent BackupProgress(Me, BackupProgressEventArgs)
    End Sub

    Public Sub CancelBackup()
        RaiseBackupEvent(BackupEvent.Cancel)
    End Sub

#Region "IDisposable Support"
    Protected Overridable Sub Dispose(ByVal disposing As Boolean)
        If Not Me.Disposed Then
            If disposing Then
                RemoveHandler Backup.BackupProgress, AddressOf ProgressManager
                If Me.BackupSchedule IsNot Nothing Then
                    Me.BackupSchedule.BackupJob = Nothing
                    Me.BackupSchedule.BackupJobSchedule = Nothing
                    Me.BackupSchedule.BackupDevice = Nothing
                    Me.BackupSchedule.VirtualMachine = Nothing
                    Me.BackupSchedule.HyperVHost = Nothing
                End If
            End If

        End If
        Me.Disposed = True
    End Sub

    ' TODO: override Finalize() only if Dispose(ByVal disposing As Boolean) above has code to free unmanaged resources.
    'Protected Overrides Sub Finalize()
    '    ' Do not change this code.  Put cleanup code in Dispose(ByVal disposing As Boolean) above.
    '    Dispose(False)
    '    MyBase.Finalize()
    'End Sub

    Public Sub Dispose() Implements IDisposable.Dispose
        Dispose(True)
        GC.SuppressFinalize(Me)
    End Sub
#End Region

End Class



Public Class BackupItem
    Public Property FilePath As String
    Public Property FileName As String
    Public Property FileSize As Long
    Public Property SnapshotDeviceObject As String
    Public Property Volume As String
    Public Property SnapshotGuid As Guid
    Public Property PathFromVolume As String

    Public Sub New()
        MyBase.New()
    End Sub

    Public Sub New(ByVal FilePath As String, ByVal FileName As String)
        Me.FilePath = FilePath
        Me.FileName = FileName
        Me.Volume = IO.Path.GetPathRoot(FilePath)
        Me.PathFromVolume = FilePath.Substring(Volume.Length)
    End Sub

End Class

Public Class BackupItemCollection
    Inherits System.Collections.ObjectModel.Collection(Of BackupItem)
End Class

Public Class VssComponentCollection
    Inherits System.Collections.ObjectModel.Collection(Of IVssWMComponent)
End Class